# analysis/go_analysis.py

from goatools.base import get_godag
from goatools.anno.gaf_reader import GafReader
import os
from goatools.go_enrichment import GOEnrichmentStudy


def run_go_enrichment(community, gaf_file="data/goa_yeast.gaf", obo_file="data/go-basic.obo"):
    """
    对给定的社区进行 GO 富集分析。
    :param community: Entrez ID 列表，代表一个社区
    :param gaf_file: GAF 格式文件路径（包含基因与 GO 的映射）
    :param obo_file: GO 本体文件路径
    :return: 富集后的 GO terms 字典 {go_term: count}
    """
    if not community:
        print("社区为空，无法进行富集分析。")
        return {}

    godag = get_godag(obo_file)

    # 读取 GAF 注释数据
    reader = GafReader(gaf_file)
    assoc = reader.get_id2gos()

    # 创建背景所有基因集合（可选：使用 GAF 文件中的所有基因）
    all_ids = set(assoc.keys())
    study_ids = set(str(g) for g in community) & all_ids

    if not study_ids:
        print("社区中没有找到有效的基因注释。")
        return {}

    # 设置富集分析参数
    goea_obj = GOEnrichmentStudy(
        all_ids,
        assoc,
        godag,
        methods=['bonferroni', 'fdr_bh']
    )

    # 执行富集分析
    results = goea_obj.run_study(study_ids)
    enriched_terms = {}

    for res in results:
        if res.p_fdr_bh < 0.05:
            enriched_terms[res.GO] = {
                'name': godag[res.GO].name,
                'pvalue': res.p_fdr_bh,
                'term_type': godag[res.GO].namespace,
                'overlap': res.ratio_in_study[0],
                'total': res.ratio_in_study[1]
            }

    return enriched_terms


def save_go_report(enriched_terms, community_id, output_dir="results/go_enrichment"):
    """
    保存 GO 富集分析报告到指定目录。
    :param enriched_terms: 富集后的 GO term 字典
    :param community_id: 社区编号（用于命名输出文件）
    :param output_dir: 输出目录
    """
    os.makedirs(output_dir, exist_ok=True)
    report_path = os.path.join(output_dir, f"go_enrichment_{community_id}.txt")

    with open(report_path, "w") as f:
        f.write(f"Community Members: {community_id}\n")
        f.write("Enriched GO Terms:\n")
        for go_id, info in enriched_terms.items():
            f.write(f"{go_id} ~ {info['name']} (p={info['pvalue']:.4f}, {info['term_type']})\n")
            f.write(f"Overlap: {info['overlap']}/{info['total']}\n")


if __name__ == "__main__":
    test_community = [813, 1234, 967]
    go_results = run_go_enrichment(test_community)
    save_go_report(go_results, "test_001")